"""
工具函数

这个模块包含与黑色素瘤检测项目相关的辅助函数。
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import os

def load_model(model_path):
    """
    加载已训练的模型
    
    参数:
        model_path: 模型文件路径
    返回:
        model: 加载的模型
    """
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
    model = tf.keras.models.load_model(model_path)
    print(f"模型已加载: {model_path}")
    return model

def preprocess_image(image, target_size=(128, 128)):
    """
    预处理输入图像以供模型使用
    
    参数:
        image: 输入图像(可以是文件路径或已加载的图像)
        target_size: 目标大小(高度, 宽度)
    返回:
        processed_img: 预处理后的图像，形状为(1, height, width, 3)
    """
    # 如果输入是字符串(文件路径)，则加载图像
    if isinstance(image, str):
        image = cv2.imread(image)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为RGB
    
    # 如果是BGR格式(OpenCV默认)，转换为RGB
    if len(image.shape) == 3 and image.shape[2] == 3:
        # 检测是否需要从BGR转换为RGB
        if isinstance(image, np.ndarray) and image.dtype == np.uint8:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 调整大小
    resized_img = cv2.resize(image, target_size)
    
    # 归一化
    normalized_img = resized_img / 255.0
    
    # 确保是float32类型
    normalized_img = normalized_img.astype(np.float32)
    
    # 添加批次维度
    batch_img = np.expand_dims(normalized_img, axis=0)
    
    return batch_img

def predict_mask(model, image):
    """
    使用模型预测分割掩码
    
    参数:
        model: 已加载的模型
        image: 原始输入图像或文件路径
    返回:
        original_image: 原始图像(RGB格式)
        mask: 预测的掩码（已调整为原始图像大小）
    """
    # 如果输入是字符串(文件路径)，则加载图像
    if isinstance(image, str):
        original_image = cv2.imread(image)
        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)  # 转换为RGB
    else:
        original_image = image.copy()
        if len(original_image.shape) == 3 and original_image.shape[2] == 3:
            # 检测是否需要从BGR转换为RGB
            if isinstance(original_image, np.ndarray) and original_image.dtype == np.uint8:
                original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
    
    # 保存原始图像尺寸以便后续恢复
    original_height, original_width = original_image.shape[:2]
    
    # 预处理图像
    processed_img = preprocess_image(original_image)
    
    # 预测掩码
    pred_mask = model.predict(processed_img)
    
    # 去除批次维度
    pred_mask = pred_mask[0]
    
    # 阈值化
    binary_mask = (pred_mask > 0.5).astype(np.uint8)
    
    # 将掩码调整回原始图像大小
    resized_binary_mask = cv2.resize(binary_mask, (original_width, original_height), 
                                    interpolation=cv2.INTER_NEAREST)
    
    return original_image, resized_binary_mask

def overlay_mask(image, mask, alpha=0.5, color=[1, 0, 0]):
    """
    将掩码叠加在原始图像上
    
    参数:
        image: 原始图像(RGB格式)
        mask: 二值掩码
        alpha: 透明度
        color: 掩码颜色 [R, G, B]
    返回:
        overlay_image: 叠加掩码的图像
    """
    # 确保图像是RGB格式
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    
    # 确保掩码与图像尺寸一致
    img_height, img_width = image.shape[:2]
    if mask.shape[:2] != image.shape[:2]:
        mask = cv2.resize(mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST)
    
    # 创建彩色掩码
    color_mask = np.zeros_like(image, dtype=np.float32)
    
    # 如果掩码是2D的，扩展为3D
    if len(mask.shape) == 2:
        mask_3d = np.expand_dims(mask, axis=-1)
    else:
        mask_3d = mask
        # 确保掩码有正确的通道数
        if mask_3d.shape[2] > 1:
            mask_3d = mask_3d[:, :, 0:1]
    
    # 设置掩码颜色
    for i in range(3):
        color_mask[:, :, i] = color[i] * mask_3d[:, :, 0]
    
    # 叠加掩码和图像
    overlay_image = image * (1 - alpha * mask_3d) + alpha * color_mask * 255
    
    # 确保值在有效范围内
    overlay_image = np.clip(overlay_image, 0, 255).astype(np.uint8)
    
    return overlay_image

def display_results(original_image, predicted_mask, true_mask=None, save_path=None):
    """
    显示原始图像、预测掩码和真实掩码(如果有)
    
    参数:
        original_image: 原始图像
        predicted_mask: 预测掩码
        true_mask: 真实掩码(可选)
        save_path: 保存路径(可选)
    """
    if true_mask is not None:
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # 显示原始图像
        axes[0].imshow(original_image)
        axes[0].set_title('原始图像')
        axes[0].axis('off')
        
        # 显示预测掩码
        axes[1].imshow(overlay_mask(original_image, predicted_mask))
        axes[1].set_title('预测掩码')
        axes[1].axis('off')
        
        # 显示真实掩码
        axes[2].imshow(overlay_mask(original_image, true_mask))
        axes[2].set_title('真实掩码')
        axes[2].axis('off')
    else:
        fig, axes = plt.subplots(1, 2, figsize=(10, 5))
        
        # 显示原始图像
        axes[0].imshow(original_image)
        axes[0].set_title('原始图像')
        axes[0].axis('off')
        
        # 显示预测掩码
        axes[1].imshow(overlay_mask(original_image, predicted_mask))
        axes[1].set_title('预测掩码')
        axes[1].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        print(f"结果已保存至 {save_path}")
    
    plt.show()
